import base64
import binascii
import json
import rsa
from Crypto.Cipher import AES
from Crypto import Random
from binascii import b2a_hex,a2b_hex
from gmssl.sm4 import CryptSM4


# 汇总各类加解密算法调用示例rsa、aes、sm4
# 1.非对称加解密rsa实现
class rsa_algorithm():
    @classmethod
    def create_rsa_keys(cls,bit_num=1024):
        public_key,private_key = rsa.newkeys(bit_num)

        public_pem = public_key.save_pkcs1()
        with open('public.pem','wb+') as f:
            f.write(public_pem)

        private_pem = private_key.save_pkcs1()
        with open('private.pem','wb+') as f:
            f.write(private_pem)
    @classmethod
    def encrypt_with_rsa(cls,plain_data):
        with open('public.pem','rb') as publickey_file:
            p = publickey_file.read()

        public_key = rsa.PublicKey.load_pkcs1(p)
        plain_data = plain_data.encode('utf-8')
        cipher_data = rsa.encrypt(plain_data,public_key)
        # print(cipher_data.decode('utf-8','ignore'))
        return cipher_data
    @classmethod
    def decrypt_with_rsa(cls,cipher_data):
        with open('private.pem','rb') as privatekey_file:
            p = privatekey_file.read()
        private_key = rsa.PrivateKey.load_pkcs1(p)
        plain_data = rsa.decrypt(cipher_data,private_key).decode()
        # print(plain_data)
        return plain_data

# 2.对称加密aes实现
# 加密数据支持字符串，不支持int类型
# 注意密钥key长度16bytes
class aes_algorithm():

    def __init__(self,key):
        self.mode = AES.MODE_CBC
        self.key = key.encode('utf-8')
        self.bs = AES.key_size[0] # 密钥长度16 bytes
        self.iv = Random.new().read(AES.key_size[0])
        self.padding = lambda s: s + (self.bs - len(s) % self.bs) * chr(self.bs - len(s) % self.bs)

    def encrypt_with_aes(self,data):
        cryptor = AES.new(key=self.key,mode=self.mode,iv=self.iv)
        base64_data = base64.b64encode(data.encode('utf-8'))
        # 转成string类型，lambda函数中s需要传入字符串类型
        str_data = str(base64_data.decode('utf-8'))
        # padding填充后加密
        cipher_data = cryptor.encrypt(self.padding(str_data).encode('utf-8'))
        # print(type(cipher_data))

        # 返回密文和原始明文长度
        return cipher_data,len(str_data)

    def decrypt_with_aes(self,cipher_data):
        # 解密过程与加密一一对应
        decryptor = AES.new(self.key,mode=self.mode,iv=self.iv)
        # 1.获取密文值进行解密，密文为string类型
        decrypted_data = decryptor.decrypt(cipher_data[0]).decode('utf-8')
        bytes_data = decrypted_data.encode('utf-8')
        base64_data = bytes_data[:cipher_data[1]]
        # 对解密后的bytes进行base64解码处理
        plain_data = base64.b64decode(base64_data).decode('utf-8')

        return plain_data

    def encrypt_with_aes_gcm(self,data):
        # aes_gcm附加数据
        associatedData = '我是中国人,this is add data'
        cipher = AES.new(key=self.key,mode=AES.MODE_GCM,nonce=self.iv)
        cipher.update(associatedData.encode())
        # aes_gcm encrypt_and_digest返回元组(密文,消息认证码MAC值)
        cipher_text,auth_tag = cipher.encrypt_and_digest(data.encode('utf-8'))
        return cipher_text,auth_tag

    def decrypt_with_aes_gcm(self,cipher_data):
        associatedData = '我是中国人,this is add data'
        cipher = AES.new(key=self.key,mode=AES.MODE_GCM,nonce=self.iv)
        cipher.update(associatedData.encode())
        plain_text = cipher.decrypt_and_verify(cipher_data[0],cipher_data[1]).decode('utf-8')
        return plain_text



# 3.国密对称加密sm4实现
class sm4_algorithm():
    pass

# if __name__ == '__main__':

    # 初始化，传入密钥key值,长度16 bytes
    # aes_en = aes_algorithm('1234567890abcdef')
    #
    # encrypted_data = aes_en.encrypt_with_aes('中国人')
    # decrypted_data = aes_en.decrypt_with_aes(encrypted_data)
    # print(decrypted_data)

    # encrypted_data = aes_en.encrypt_with_aes_gcm('我是中国人')
    # print('en:aes-gcm:',encrypted_data)
    # decrypted_data = aes_en.decrypt_with_aes_gcm(encrypted_data)
    # print('de:aes-gcm:',decrypted_data)



    # rsa_algorithm.create_rsa_keys()
    # cipher = rsa_algorithm.encrypt_with_rsa('我是中国人')
    # plain = rsa_algorithm.decrypt_with_rsa(cipher)
    # print("cipher>>>: ",cipher,'\n',"plain>>>: ",plain)


